#!/usr/bin/env python
"""

    run_gencompress_size.py

"""

################################################################################
#
#   run_gencompress_size
#
#
#   Copyright (C) 2009 Leo Goodstadt
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation; version 2
#   of the License
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#################################################################################

import sys, os
from collections import defaultdict
import json
#import popen2
import subprocess
import tempfile
import string
import time

#
#   translation table for filtering out invalid DNA
#
_valid_dna = string.ascii_uppercase        #"ACGTURYKMSWBDHVN"
_trans_valid_dna = string.maketrans(_valid_dna.lower(), _valid_dna)
_allchars = string.maketrans('','')
_delete_invalid_dna = _allchars.translate(_allchars, _valid_dna.translate(_allchars, ""))




#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#_________________________________________________________________________________________
#
#   run_gencompress_size
#
#_________________________________________________________________________________________
def _write_temp_dna_file (dna_sequence):
    """
    returns file name of written sequence
    """
    # write dna file
    dna_sequence_file = tempfile.NamedTemporaryFile(suffix='.gencompress', delete=False)
    dna_sequence_file_name = str(dna_sequence_file.name)
    written_dna_sequence = dna_sequence.translate(_trans_valid_dna, _delete_invalid_dna)
    dna_sequence_file.write(written_dna_sequence)
    dna_sequence_file.flush()
    dna_sequence_file.close()
    del dna_sequence_file
    return dna_sequence_file_name

def _cleanup_files (file_name1, file_name2):
    """
    Cleanup temp files
    """
    file_name1_root = file_name1[0:file_name1.find(".")]
    try:
        os.unlink(file_name1)
        os.unlink(file_name1_root + ".GEN")
        os.unlink(file_name1_root + ".LOG")
        if file_name2:
            os.unlink(file_name2)
    except:
        pass


def call_gencompress (cmds):
    """
    makes actual call to gencompress.
    Abstracted so we can call this twice if the first throws
    """
    for i in range(5):
        process = subprocess.Popen(cmds,
                                   stdin=subprocess.PIPE,
                                   stderr=subprocess.PIPE,
                                   stdout=subprocess.PIPE,
                                   shell=True
                                   )
        (stdout_str, stderr_str) = process.communicate()
        if process.returncode == 0:
            return 0, stdout_str, stderr_str
        time.sleep(1)
    return process.returncode, stdout_str, stderr_str


def run_gencompress_size (dna_sequence, gencompress_bin = None, dna_sequence2 = None):
    """
    Calculates the compressed size of sequences using the GenCompress programme
        http://www1.spms.ntu.edu.sg/~chenxin/GenCompress/

    Returns the size and compressed sizes as a tuple

    If dna_sequence2 is given, calculates the conditional compression with dna_sequence2
        as additional input

    N.B. Non-alphabetic letters are ignored, and all sequence is converted to uppercase
    """
    if not gencompress_bin:
        gencompress_bin = "GenCompress"

    dna_sequence_file_name = _write_temp_dna_file (dna_sequence)

    aux_sequence_file_name = None
    aux_sequence_cmd_str = ""
    if dna_sequence2:
        aux_sequence_file_name = _write_temp_dna_file (dna_sequence2)
        aux_sequence_cmd_str = "-c %s" % aux_sequence_file_name


    try:
        cmds = " ".join((gencompress_bin, dna_sequence_file_name, aux_sequence_cmd_str))
        returncode, stdout_str, stderr_str = call_gencompress(cmds)
        if returncode == 0:
            compressed_file_name = dna_sequence_file_name[0:dna_sequence_file_name.find(".")] + ".GEN"
            if not os.path.exists(compressed_file_name):
                err_str = ("GenCompress failed to compress for [%s]\n%s\n%s\n" % (
                                                                        dna_sequence_file_name,
                                                                        stdout_str,
                                                                        stderr_str))
                raise Exception(err_str)
            compressed_size = os.path.getsize(compressed_file_name)

            _cleanup_files (dna_sequence_file_name, aux_sequence_file_name)
            return len(dna_sequence), compressed_size

        #
        #   failed: cleanup and throw exception
        #
        else:

            err_str = ("GenCompress failed, returning "
                        "%s for [%s]\n[%s %s %s]\n%s\n%s\n" % (returncode,
                                                               dna_sequence,
                                                               cmds,
                                                               dna_sequence_file_name,
                                                               aux_sequence_file_name,
                                                               stdout_str,
                                                               stderr_str))
            raise Exception(err_str)
    except:
        #_cleanup_files (dna_sequence_file_name, aux_sequence_file_name)
        raise



#_________________________________________________________________________________________
#
#   pairwise_all_vs_all_gencompress_size
#
#_________________________________________________________________________________________
def pairwise_all_vs_all_gencompress_size (dna_sequences, gencompress_bin = None):
    """
    Returns square matrix of (original length, compressed length) tuples
    Diagonals are self-self compressions
    Pairwise compressions are concatenations
    """
    cnt_seq = len(dna_sequences)
    results = [[None] * cnt_seq for jj in range(cnt_seq)]
    for ii in range(cnt_seq):
        for jj in range(cnt_seq):
            if ii == jj:
                sys.stderr.write("1")
                results[ii][ii] = run_gencompress_size(dna_sequences[ii])
            else:
                sys.stderr.write("1")
                results[ii][jj] = run_gencompress_size(dna_sequences[ii] + dna_sequences[jj])
    return results

#_________________________________________________________________________________________
#
#   conditional_pairwise_all_vs_all_gencompress_size
#
#_________________________________________________________________________________________
def conditional_pairwise_all_vs_all_gencompress_size (dna_sequences, gencompress_bin = None):
    """
    Returns square matrix of (original length, compressed length) tuples
    Diagonals are self-self compressions
    Pairwise compressions are conditional compressions
    """
    cnt_seq = len(dna_sequences)
    results = [[None] * cnt_seq for jj in range(cnt_seq)]
    for ii in range(cnt_seq):
        for jj in range(cnt_seq):
            if ii == jj:
                sys.stderr.write("2")
                results[ii][ii] = run_gencompress_size(dna_sequences[ii])
            else:
                sys.stderr.write("2")
                results[ii][jj] = run_gencompress_size(dna_sequences[ii], dna_sequence2 = dna_sequences[jj])
    return results



def pairwise_compression_matrix_to_ncd (compression_matrix):
    """
    calculate ncd = (min(size_ab, size_ba) - min(size_a, size_b)) / max(size_a, size_b)
    """
    # alias with shorter name
    m = compression_matrix
    cnt_seq = len(m)
    results = [[100] * cnt_seq for jj in range(cnt_seq)]


    for a in range(cnt_seq):
        for b in range(a):
            results[a][b] = ( (min(m[a][b][1], m[b][a][1]) - min(m[a][a][1], m[b][b][1])) * 1.0 /
                                    max(m[a][a][1], m[b][b][1]))
            results[b][a] = results[a][b]

    return results

#
def pairwise_compression_matrix_to_ucd (compression_matrix):
    """
    calculate ucd = max(size_ab - size_a, size_ba - size_b) / max(size_a, size_b)
    """
    # alias with shorter name
    m = compression_matrix
    cnt_seq = len(m)
    results = [[100] * cnt_seq for jj in range(cnt_seq)]


    for a in range(cnt_seq):
        for b in range(a):
            results[a][b] = ( max(m[a][b][1] - m[a][a][1], m[b][a][1] - m[b][b][1]) * 1.0 /
                                    max(m[a][a][1], m[b][b][1]))
            results[b][a] = results[a][b]

    return results

#
def pairwise_conditional_compression_matrix_to_usm (compression_matrix):
    """
    calculate usm = max(C(x|y), C(y|x)) / max(C(x), C(y))
    """
    # alias with shorter name
    m = compression_matrix
    cnt_seq = len(m)
    results = [[100] * cnt_seq for jj in range(cnt_seq)]


    for a in range(cnt_seq):
        for b in range(a):
            results[a][b] = ( max(m[a][b][1], m[b][a][1]) * 1.0 /
                                    max(m[a][a][1], m[b][b][1]))
            results[b][a] = results[a][b]

    return results



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
import unittest
class Test_run_gencompress_size(unittest.TestCase):

    #       self.assertEqual(self.seq, range(10))
    #       self.assert_(element in self.seq)
    #       self.assertRaises(ValueError, random.sample, self.seq, 20)



    def test_function(self):
        """
            test
        """
        import re
        from collections import defaultdict

        base_dir = "test/test_run_gencompress_size/"


        for seq_file_name in "cdna", "prot", "cdna_prot":
            sequences_by_name = defaultdict(str)
            seq_file = base_dir + seq_file_name + ".fa"
            curr_name = None
            for line in open(seq_file):
                line = line.rstrip()
                if not len(line):
                    continue
                if line[0] == '>':
                    curr_name = "_".join(re.split("_+", line[1:])[:-1])
                else:
                    sequences_by_name[curr_name] += line


            for name, s in sorted(sequences_by_name.iteritems()):
                orig_size, compressed_size = run_gencompress_size(s)
                print "%50s %10d %10d %10.1f%%" % (name, orig_size, compressed_size, compressed_size * 100.0 / orig_size)

            names = sorted(sequences_by_name.keys())
            sequences = [sequences_by_name[n] for n in names]

            compressed_matrix = pairwise_all_vs_all_gencompress_size (sequences)
            ncd_matrix = pairwise_compression_matrix_to_ncd (compressed_matrix)
            ucd_matrix = pairwise_compression_matrix_to_ucd (compressed_matrix)

            compressed_matrix = conditional_pairwise_all_vs_all_gencompress_size (sequences)
            usm_matrix = pairwise_conditional_compression_matrix_to_usm(compressed_matrix)


            cnt_seq = len(names)
            results_file = open("test/test_run_gencompress_size/ncd.%s.pairs" % seq_file_name, "w")
            for ii in range(cnt_seq):
                for jj in range(ii):
                    print >>results_file, "%s\t%s\t%g" % (names[ii], names[jj], ncd_matrix[ii][jj])
            results_file = open("test/test_run_gencompress_size/ucd.%s.pairs" % seq_file_name, "w")
            for ii in range(cnt_seq):
                for jj in range(ii):
                    print >>results_file, "%s\t%s\t%g" % (names[ii], names[jj], ucd_matrix[ii][jj])

            results_file = open("test/test_run_gencompress_size/usm.%s.pairs" % seq_file_name, "w")
            for ii in range(cnt_seq):
                for jj in range(ii):
                    print >>results_file, "%s\t%s\t%g" % (names[ii], names[jj], usm_matrix[ii][jj])



#
#   debug code not run if called as a module
#
if __name__ == '__main__':
    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
    unittest.main()




